# Thompson sampling. 
# Based on https://arxiv.org/abs/1111.1797. 

mutable struct ThompsonSampling <: Policy end

struct ThompsonSamplingDetails <: PolicyDetails
  solverTime::Float64
end

function choose_action(instance::CombinatorialInstance, policy::ThompsonSampling, state::State{T}; with_trace::Bool=false) where T
  # Determine the weights for each arm and use them to solve the combinatorial problem. 
  # println("Counts: $(state.arm_counts); rewards: $(state.arm_average_reward)")
  weights = Dict(arm => rand(Beta(state.arm_counts[arm] * state.arm_average_reward[arm] + 1, 
                                  state.arm_counts[arm] * (1 - state.arm_average_reward[arm]) + 1)) for arm in keys(state.arm_counts))
  
  t0 = now()
  sol = solve_linear(instance, weights)
  t1 = now()

  if with_trace
    return sol, ThompsonSamplingDetails((t1 - t0).value)
  else
    return sol
  end
end
